#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
extern void onehot_custom_do(uint32_t blockDim, void *stream, uint8_t *x,uint8_t *depth, uint8_t *z);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void onehot_custom(GM_ADDR x,GM_ADDR depth, GM_ADDR z);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    size_t inputByteSize_x = 8*1024*sizeof(int32_t);
    size_t inputByteSize_depth = 63*sizeof(int32_t);
    size_t outputByteSize = 63*8*1024*sizeof(int32_t);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputByteSize_x);
    uint8_t *depth = (uint8_t *)AscendC::GmAlloc(inputByteSize_depth);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    ReadFile("./input/input_x.bin", inputByteSize_x, x, inputByteSize_x);
    ReadFile("./input/input_depth.bin", inputByteSize_depth, depth, inputByteSize_depth);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(onehot_custom, blockDim, x,depth, z); // use this macro for cpu debug

    WriteFile("./output/output_z.bin", z, outputByteSize);

    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)depth);
    AscendC::GmFree((void *)z);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *xHost,*depthHost,*zHost;
    uint8_t *xDevice,*depthDevice,*zDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputByteSize_x));
    CHECK_ACL(aclrtMallocHost((void **)(&depthHost), inputByteSize_depth));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost), outputByteSize));

    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize_x, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&depthDevice, inputByteSize_depth, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_x.bin", inputByteSize_x, xHost, inputByteSize_x);
    ReadFile("./input/input_depth.bin", inputByteSize_depth, depthHost, inputByteSize_depth);

    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize_x, xHost, inputByteSize_x, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(depthDevice, inputByteSize_depth, depthHost, inputByteSize_depth, ACL_MEMCPY_HOST_TO_DEVICE));

    onehot_custom_do(blockDim, stream, xDevice,depthDevice,zDevice);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, outputByteSize);

    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(depthDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(depthHost));
    CHECK_ACL(aclrtFreeHost(zHost));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}
